from ENF_deepLearning.utils.cnn.Net import Net
from ENF_deepLearning.utils.cnn.Net_Dataset import Net_Dataset
from ENF_deepLearning.utils.parent.Base import Base
import torch
from torch.utils.data import DataLoader

import numpy as np
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class CNN_Classifier(Base):
    def one_D_cnn(self, x_train, x_test, y_train, y_test,    # 输入维度(sample_size, 600)，
                  epoch=10, batch_size=16, lr=0.001, print_for_epoch=10,
                  n_filter=65, kernel_size=60, stride=1, number_class=10,filter_row=2,
                  loss_function=torch.nn.CrossEntropyLoss, optimizer=torch.optim.SGD):

        print('x_train shape:', x_train.shape, 'x_test shape', x_test.shape)

        net = Net(n_filter=n_filter, kernel_size=kernel_size, stride=stride, number_class=number_class, filter_row=filter_row)
        net = net.to(device)
        criterion = loss_function()
        optimizer = optimizer(net.parameters(), lr=lr, momentum=0.9)
        ss = Net_Dataset(x_train, y_train)
        train_loader = DataLoader(dataset=ss, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)

        mean_acc = []
        test_acc = []

        x_test = np.array(x_test, dtype=np.float)
        y_test = np.array(y_test, dtype=np.longlong)
        xx = torch.from_numpy(x_test).to(device)
        yy = torch.from_numpy(y_test).to(device)
        correct = 0
        for it in range(epoch):
            loss_sum = 0
            count = 0
            for index, (X, y) in enumerate(train_loader):
                X = X.to(device)
                y = y.to(device)
                predict = net(X)
                loss = criterion(predict, y)
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                loss_sum += loss.item()
                count += 1
            predict = net(xx)
            prediction = torch.argmax(predict, dim=1) == yy
            acc = (prediction.sum().float().cpu().numpy() / x_test.shape[0])# 测试集的准确率
            correct += acc
            mean_acc.append(correct / (it + 1))
            test_acc.append(acc)
            # print(acc)
            if (it + 1) % print_for_epoch == 0:
                print('epoch:', it + 1, 'loss:', loss_sum / count, 'test acc:', correct / (it + 1), 'cur acc:', acc)
        predict = net(xx)
        y_pred = torch.argmax(predict, dim=1).int().cpu().numpy().tolist()
        y_true = yy.int().cpu().numpy().tolist()

        self.save_pred_true(y_pred, y_true)

        print('max test acc:', np.max(np.array(test_acc)))
        return np.max(test_acc)
        # plt.subplot(211)
        # plt.plot(mean_acc)
        # plt.title('mean acc')
        # plt.subplot(212)
        # plt.plot(test_acc)
        # plt.title('test acc for per epoch')
        # plt.tight_layout()
        # plt.show()